package utils

import (
	"fmt"
	"time"
)

type RateLimit struct {
	list   *Queue
	opened bool
	size   uint64
	sec    uint64
}

// NewRateLimit 获得一个实例，每 sec 秒 可以处理 size 次
func NewRateLimit(size, sec uint64) (rl *RateLimit, err error) {
	var lst *Queue
	lst, err = NewQueue(size)
	if err != nil {
		return
	}
	if sec <= 0 {
		err = fmt.Errorf("\"sec\" 必须大于0")
		return
	}
	rl = &RateLimit{
		list:   lst,
		opened: false,
		size:   size,
		sec:    sec,
	}
	rl.Open()
	return
}

// Open 启动
func (rl *RateLimit) Open() {
	if rl.opened {
		return
	}
	rl.opened = true
	rl.list.Pause(false)
	go func() {
		for td := time.Duration(rl.sec*1000000/rl.size) * time.Microsecond; rl.opened; {
			rl.list.Push(struct{}{})
			time.Sleep(td)
		}
	}()
}

// Close 关闭
func (rl *RateLimit) Close() {
	if !rl.opened {
		return
	}
	rl.list.Pause(true)
	rl.opened = false
}

// Status 当前可用量
func (rl *RateLimit) Status() int {
	return len(rl.list.data)
}

// Check 是否可用，非阻塞
func (rl *RateLimit) Check() bool {
	if !rl.opened {
		return true
	}
	_, ok := rl.list.Pop()
	return ok
}

// CheckWait 等待可用，阻塞，可设置超时
func (rl *RateLimit) CheckWait(td time.Duration) bool {
	if !rl.opened {
		return true
	}
	_, err := rl.list.PopWait(td)
	return err == nil
}
